import torch
import torch.nn as nn
import torch.nn.functional as F

class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.InstanceNorm3d(out_channels),
            nn.LeakyReLU(0.2, inplace=True),  # 使用LeakyReLU保留负值信息
            nn.Dropout3d(p=0.1),
            nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.InstanceNorm3d(out_channels),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Dropout3d(p=0.1)  # 降低dropout比例
        )
        self.residual = nn.Conv3d(in_channels, out_channels, kernel_size=1) if in_channels != out_channels else None
        
    def forward(self, x):
        identity = x
        out = self.double_conv(x)
        if self.residual is not None:
            identity = self.residual(identity)

        diff = [out.shape[i] - identity.shape[i] for i in range(2, 5)]
        identity = F.pad(identity, [0, diff[2], 0, diff[1], 0, diff[0]])
        return out + identity

class ASPP(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        rates = [1, 2, 4]
        
        self.convs = nn.ModuleList()
        for rate in rates:
            self.convs.append(
                nn.Sequential(
                    nn.Conv3d(in_channels, out_channels, kernel_size=3, 
                              padding=rate, dilation=rate),
                    nn.InstanceNorm3d(out_channels),
                    nn.LeakyReLU(0.2, inplace=True)
                )
            )
        
        self.global_conv = nn.Sequential(
            nn.Conv3d(in_channels, out_channels, kernel_size=1),
            nn.LeakyReLU(0.2, inplace=True)
        )
        
        self.project = nn.Sequential(
            nn.Conv3d(out_channels * (len(rates) + 1), out_channels, kernel_size=1),
            nn.InstanceNorm3d(out_channels),
            nn.LeakyReLU(0.2, inplace=True)
        )
        
        self.avgpool = nn.AdaptiveAvgPool3d(1)
        
    def forward(self, x):
        features = [conv(x) for conv in self.convs]
        
        global_feat = self.avgpool(x)
        global_feat = self.global_conv(global_feat)
        
        global_feat = F.interpolate(global_feat, size=x.shape[2:], mode='trilinear', align_corners=True)
        features.append(global_feat)
        
        return self.project(torch.cat(features, dim=1))

class Down(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool3d(2),
            DoubleConv(in_channels, out_channels)
        )
        
    def forward(self, x):
        return self.maxpool_conv(x)

class Up(nn.Module):
    def __init__(self, in_channels, out_channels, use_attn=False):
        super().__init__()
        self.up = nn.ConvTranspose3d(in_channels, in_channels // 2, kernel_size=2, stride=2)
        if use_attn:
            self.attn_gate = nn.Sequential(
                nn.Conv3d(in_channels // 2, in_channels // 2, kernel_size=1),
                nn.InstanceNorm3d(in_channels // 2),
                nn.Sigmoid()
            )
        else:
            self.attn_gate = None
            
        self.conv = DoubleConv(in_channels, out_channels)
        
    def forward(self, x1, x2):
        x1 = self.up(x1)
        
        if self.attn_gate is not None:
            attention = self.attn_gate(x2)
            x2 = x2 * attention
        
        diff = [x2.shape[i] - x1.shape[i] for i in range(2, 5)]
        x1 = F.pad(x1, [0, diff[2], 0, diff[1], 0, diff[0]])
        
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)

class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size=1)
        
    def forward(self, x):
        return self.conv(x)

class AuxiliaryHead(nn.Module):  # 深度监督辅助输出
    def __init__(self, in_channels, out_channels, scale_factor):
        super().__init__()
        self.conv = OutConv(in_channels, out_channels)
        self.scale_factor = scale_factor
        
    def forward(self, x):
        x = self.conv(x)
        return F.interpolate(x, scale_factor=self.scale_factor, mode='trilinear', align_corners=True)

class UNet3D_ASPP(nn.Module):
    def __init__(self, n_channels=1, n_classes=1):
        super(UNet3D_ASPP, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        
        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        self.down4 = DoubleConv(512, 1024)  # 不再下采样，保留更多空间信息
        
        self.aspp = ASPP(1024, 1024)
        
        self.up1 = Up(1024, 512, use_attn=True)
        self.up2 = Up(512, 256, use_attn=True)
        self.up3 = Up(256, 128, use_attn=True)
        self.up4 = Up(128, 64, use_attn=True)
        
        self.outc = OutConv(64, n_classes)
        
        self.aux_out1 = AuxiliaryHead(512, n_classes, scale_factor=8)
        self.aux_out2 = AuxiliaryHead(256, n_classes, scale_factor=4)
        
    def forward(self, x):
        x1 = self.inc(x)        # -> [B, 64, 64, 64, 64]
        x2 = self.down1(x1)     # -> [B, 128, 32, 32, 32]
        x3 = self.down2(x2)     # -> [B, 256, 16, 16, 16]
        x4 = self.down3(x3)     # -> [B, 512, 8, 8, 8]
        x5 = self.down4(x4)     # -> [B, 1024, 8, 8, 8]
        
        x5 = self.aspp(x5)      # -> [B, 1024, 8, 8, 8]
        
        x = self.up1(x5, x4)    # -> [B, 512, 16, 16, 16]
        aux1 = self.aux_out1(x) # 辅助输出1
        
        x = self.up2(x, x3)     # -> [B, 256, 32, 32, 32]
        aux2 = self.aux_out2(x) # 辅助输出2
        
        x = self.up3(x, x2)     # -> [B, 128, 64, 64, 64]
        x = self.up4(x, x1)     # -> [B, 64, 64, 64, 64]
        
        logits = self.outc(x)   # -> [B, 1, 64, 64, 64]
        
        return logits

if __name__ == "__main__":
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")
    if torch.cuda.is_available():
        print(f"GPU型号: {torch.cuda.get_device_name(0)}")
    
    input_tensor = torch.randn(2, 1, 64, 64, 64).to(device)
    
    model = UNet3D_ASPP(n_channels=1, n_classes=1).to(device)
    output, aux_out1, aux_out2 = model(input_tensor)
    
    print("输入维度:", input_tensor.shape)
    print("主输出维度:", output.shape)        # [2, 1, 64, 64, 64]
    print("辅助输出1维度:", aux_out1.shape)   # [2, 1, 64, 64, 64]
    print("辅助输出2维度:", aux_out2.shape)   # [2, 1, 64, 64, 64]

    total_params = sum(p.numel() for p in model.parameters())
    print(f"模型总参数量: {total_params/1e6:.2f}M")

    